﻿#pragma once

#include "gopt_base.h"
#include <chrono>
#include <iomanip>
#include <iostream>

namespace GOptSolver {

using aug_lag_evaluate_func_type = evaluate_func_type;

inline double clamp(double min, double max, double x) {
  return std::min(max, std::max(min, x));
}

struct AugLagState {
  double penalty = 0;
  bool feasible = true;
  double ICM = HUGE_VAL;
  double rho = 1.0;
  int iteration = -1;
  long long time_seconds = 0;
};

struct AugLagParam {
  double tau = 0.5;
  double gam = 10;
  double lambda_min = -1e20;
  double lambda_max = 1e20;
  int max_iterations = 0;
  int max_time_seconds = 0;
  double delta = 1e-6;
  double g_epsilon = 1e-6;
};

using print_auglag_progress_func_type =
    std::function<int(AugLagState &state, const double fx,
                      const Eigen::VectorXd &x, const Eigen::VectorXd &g)>;

inline int default_auglag_print_progress(AugLagState &state, const double fx,
                                         const Eigen::VectorXd &x,
                                         const Eigen::VectorXd &g) {
  std::cout << std::setprecision(4)
            << "=============AugLag=================" << std::endl
            << "Iteration: " << state.iteration << std::endl
            << "Penalty: " << state.penalty << std::endl
            << "Feasible: " << (state.feasible ? "ok" : "not") << std::endl
            << "Rho: " << state.rho << std::endl
            << "ICM: " << state.ICM << std::endl
            << "Gradient Inf Norm: " << g.cwiseAbs().maxCoeff() << std::endl
            << "Function Value: " << fx << std::endl
            << "Variables: " << std::endl
            << x.transpose() << std::endl
            << "Time(s): " << state.time_seconds << std::endl;
  return 0;
}

class LagConstraintBase {
public:
  virtual double update_lambda(AugLagState &state, const AugLagParam &params,
                               const Eigen::VectorXd &x,
                               Eigen::VectorXd &g) = 0;
  virtual double forward(AugLagState &state, const AugLagParam &params,
                         const Eigen::VectorXd &x, Eigen::VectorXd &g) = 0;
};

using aug_lag_normal_evaluate_func_type =
    std::function<void(void *instance, const Eigen::VectorXd &x,
                       Eigen::VectorXd &f, Eigen::MatrixXd &g)>;

// x <= u
class LagBoundContraint : public LagConstraintBase
{
private:
  Eigen::VectorXd bounds_;
  bool is_upper_bound_ = true;
private:
  Eigen::VectorXd lambda_;
  double tol_ = 1e-6;

public:
  virtual void set_constraint(Eigen::VectorXd bounds, bool is_upper_bound) {
    auto num_constraints = bounds.size();
    bounds_ = bounds;
    is_upper_bound_ = is_upper_bound;
    lambda_.setZero(num_constraints);
  }

  double update_lambda(AugLagState &state, const AugLagParam &params,
                       const Eigen::VectorXd &x, Eigen::VectorXd &g) {
    Eigen::VectorXd fci = x - bounds_;
    if (!is_upper_bound_) {
        fci *= -1.0;
    }
    Eigen::VectorXd newlambda_ = lambda_ + state.rho * fci;
    state.penalty += fci.cwiseMax(0).sum();
    state.feasible = state.feasible && (fci.array() <= tol_).all();
    state.ICM = std::max(state.ICM, 
    fci.cwiseMax(-lambda_ / state.rho).cwiseAbs().maxCoeff());
    lambda_ = newlambda_.cwiseMax(0).cwiseMin(params.lambda_max);
    return fci.cwiseMax(0).sum();
  }

  double forward(AugLagState &state, const AugLagParam &params,
                 const Eigen::VectorXd &x, Eigen::VectorXd &g) {
    Eigen::VectorXd h = lambda_ / state.rho;
    if (is_upper_bound_) {
        h += x - bounds_;
    } else {
        h += bounds_ - x;
    }
    h = h.cwiseMax(0);
    double L = 0.5 * state.rho * h.squaredNorm();
    g = state.rho * h;
    if (!is_upper_bound_) {
        g *= -1.0;
    }
    return L;
  }
};

// l <= x <= u
class LagBoxContraint : public LagConstraintBase {
private:
  Eigen::VectorXd upper_bounds_;
  Eigen::VectorXd lower_bounds_;
private:
  Eigen::VectorXd lambda_;
  Eigen::VectorXd z_;
  double alpha_ = 1.6;
  double tol_ = 1e-6;

public:
  virtual void set_constraint(Eigen::VectorXd lower_bounds, Eigen::VectorXd upper_bounds) {
    auto num_constraints = upper_bounds.size();
    lower_bounds_ = lower_bounds;
    upper_bounds_ = upper_bounds;
    lambda_.setZero(num_constraints);
    z_.setZero(num_constraints);
  }

  double update_lambda(AugLagState &state, const AugLagParam &params,
                       const Eigen::VectorXd &x, Eigen::VectorXd &g) override {
    Eigen::VectorXd f_cache = x - z_;

    Eigen::VectorXd z_next_bar = f_cache * alpha_ + z_;
    Eigen::VectorXd z_next = z_next_bar + lambda_ / state.rho;
    z_ = z_next.cwiseMin(upper_bounds_).cwiseMax(lower_bounds_);

    auto new_lambda_ = lambda_ + state.rho * (z_next_bar - z_);
    lambda_ =
        new_lambda_.cwiseMin(params.lambda_max).cwiseMax(params.lambda_min);
    f_cache = f_cache.cwiseAbs();
    state.penalty += f_cache.sum();
    for (int i = 0; i < f_cache.size(); ++i) {
      state.feasible = state.feasible && f_cache[i] <= tol_;
    }
    state.ICM = std::max(state.ICM, f_cache.maxCoeff());
    return f_cache.sum();
  }
  double forward(AugLagState &state, const AugLagParam &params,
                 const Eigen::VectorXd &x, Eigen::VectorXd &g) override {
    Eigen::VectorXd f_cache = x - z_ + lambda_ / state.rho;
    auto L = 0.5 * state.rho * f_cache.squaredNorm();
    g = state.rho * f_cache;
    return L;
  }
};

// l <= Ax + b <= u
class LagBoxFuncContraint : public LagConstraintBase {
private:
  aug_lag_normal_evaluate_func_type func_;
  void *instance_;
  Eigen::VectorXd f_cache_;
  Eigen::VectorXd lambda_;
  Eigen::VectorXd z_;
  Eigen::VectorXd lower_bounds_;
  Eigen::VectorXd upper_bounds_;
  Eigen::MatrixXd g_cache_;
  double alpha_ = 1.6;
  double tol_ = 1e-6;

public:
  virtual void set_constraint(const aug_lag_normal_evaluate_func_type &func,
                              Eigen::VectorXd lower_bounds,
                              Eigen::VectorXd upper_bounds, void *instance) {
    func_ = func;
    auto num_constraints = upper_bounds.size();
    lower_bounds_ = lower_bounds;
    upper_bounds_ = upper_bounds;
    instance_ = instance;
    f_cache_.setZero(num_constraints);
    lambda_.setZero(num_constraints);
    z_.setZero(num_constraints);
  }

  virtual void evaluate(const Eigen::VectorXd &x, Eigen::VectorXd &f,
                        Eigen::MatrixXd &g) {
    func_(instance_, x, f, g);
    f -= z_;
  }

  double update_lambda(AugLagState &state, const AugLagParam &params,
                       const Eigen::VectorXd &x, Eigen::VectorXd &g) override {
    Eigen::MatrixXd g_cache_;
    g_cache_.resize(f_cache_.size(), x.size());
    evaluate(x, f_cache_, g_cache_);

    Eigen::VectorXd z_next_bar = f_cache_ * alpha_ + z_;
    Eigen::VectorXd z_next = z_next_bar + lambda_ / state.rho;
    z_ = z_next.cwiseMin(upper_bounds_).cwiseMax(lower_bounds_);
    auto new_lambda_ = lambda_ + state.rho * (z_next_bar - z_);
    lambda_ =
        new_lambda_.cwiseMin(params.lambda_max).cwiseMax(params.lambda_min);

    f_cache_ = f_cache_.cwiseAbs();
    state.penalty += f_cache_.sum();
    for (int i = 0; i < f_cache_.size(); ++i) {
      state.feasible = state.feasible && f_cache_[i] <= tol_;
    }
    state.ICM = std::max(state.ICM, f_cache_.maxCoeff());
    return f_cache_.sum();
  }
  double forward(AugLagState &state, const AugLagParam &params,
                 const Eigen::VectorXd &x, Eigen::VectorXd &g) override {
    g_cache_.resize(f_cache_.size(), x.size());
    evaluate(x, f_cache_, g_cache_);
    f_cache_ = f_cache_ + lambda_ / state.rho;
    auto L = 0.5 * state.rho * f_cache_.squaredNorm();
    g = state.rho * (g_cache_.transpose() * f_cache_);
    return L;
  }
};

class LagConstraintNormalBase : public LagConstraintBase {
private:
  aug_lag_normal_evaluate_func_type func_;
  void *instance_;

protected:
  Eigen::VectorXd f_cache_;
  Eigen::VectorXd lambda_;
  Eigen::MatrixXd g_cache_;

public:
  virtual void set_constraint(const aug_lag_normal_evaluate_func_type &func,
                              int num_constraints, int num_variables,
                              void *instance) {
    func_ = func;
    f_cache_.setZero(num_constraints);
    lambda_.setZero(num_constraints);
    g_cache_.setZero(num_constraints, num_variables);
    instance_ = instance;
  }

  virtual void evaluate(const Eigen::VectorXd &x, Eigen::VectorXd &f,
                        Eigen::MatrixXd &g) {
    return func_(instance_, x, f, g);
  }

  virtual double update_lambda(AugLagState &state, const AugLagParam &params,
                               const Eigen::VectorXd &x,
                               Eigen::VectorXd &g) = 0;
  virtual double forward(AugLagState &state, const AugLagParam &params,
                         const Eigen::VectorXd &x, Eigen::VectorXd &g) = 0;
};

// F(x) = 0
class LagEqConstraintNormal : public LagConstraintNormalBase {
private:
  double tol_ = 1e-6;

public:
  double update_lambda(AugLagState &state, const AugLagParam &params,
                       const Eigen::VectorXd &x, Eigen::VectorXd &g) {
    evaluate(x, f_cache_, g_cache_);

    auto new_lambda_ = lambda_ + state.rho * f_cache_;
    lambda_ =
        new_lambda_.cwiseMin(params.lambda_max).cwiseMax(params.lambda_min);

    f_cache_ = f_cache_.cwiseAbs();
    state.penalty += f_cache_.sum();
    for (int i = 0; i < f_cache_.size(); ++i) {
      state.feasible = state.feasible && f_cache_[i] <= tol_;
    }
    state.ICM = std::max(state.ICM, f_cache_.maxCoeff());

    return f_cache_.sum();
  }

  virtual double forward(AugLagState &state, const AugLagParam &params,
                         const Eigen::VectorXd &x, Eigen::VectorXd &g) {
    evaluate(x, f_cache_, g_cache_);
    f_cache_ = f_cache_ + lambda_ / state.rho;
    auto L = 0.5 * state.rho * f_cache_.squaredNorm();
    g = state.rho * (g_cache_.transpose() * f_cache_);
    return L;
  }
};

// F(x) <= 0
class LagInEqConstraintNormal : public LagConstraintNormalBase {
private:
  double tol_ = 1e-6;

public:
  double update_lambda(AugLagState &state, const AugLagParam &params,
                       const Eigen::VectorXd &x, Eigen::VectorXd &g) {
    evaluate(x, f_cache_, g_cache_);
    auto newlambda_ = lambda_ + state.rho * f_cache_;
    state.ICM = std::max(
        state.ICM,
        fabs(std::max(f_cache_.maxCoeff(), -lambda_.minCoeff() / state.rho)));
    for (int i = 0; i < f_cache_.size(); ++i) {
      state.feasible = state.feasible && f_cache_[i] <= tol_;
    }
    state.penalty += f_cache_.cwiseMax(0).sum();
    lambda_ = newlambda_.cwiseMax(0).cwiseMin(params.lambda_max);
    return f_cache_.cwiseMax(0).sum();
  }

  virtual double forward(AugLagState &state, const AugLagParam &params,
                         const Eigen::VectorXd &x, Eigen::VectorXd &g) {
    evaluate(x, f_cache_, g_cache_);
    f_cache_ = f_cache_ + lambda_ / state.rho;
    f_cache_ = f_cache_.cwiseMax(0);
    auto L = 0.5 * state.rho * f_cache_.squaredNorm();
    g = state.rho * (g_cache_.transpose() * f_cache_);
    return L;
  }
};

class LagConstraintScalarBase : public LagConstraintBase {
private:
  aug_lag_evaluate_func_type func_;
  void *instance_;

public:
  virtual void set_constraint(const aug_lag_evaluate_func_type &func,
                              void *instance) {
    func_ = func;
    instance_ = instance;
  }
  virtual double evaluate(const Eigen::VectorXd &x, Eigen::VectorXd &g) {
    return func_(instance_, x, g);
  }

  virtual double update_lambda(AugLagState &state, const AugLagParam &params,
                               const Eigen::VectorXd &x,
                               Eigen::VectorXd &g) = 0;
  virtual double forward(AugLagState &state, const AugLagParam &params,
                         const Eigen::VectorXd &x, Eigen::VectorXd &g) = 0;
};

// f(x) = 0
class LagEqConstraint : public LagConstraintScalarBase {
private:
  double tol_ = 1e-6;
  double lambda_ = 0;

public:
  double update_lambda(AugLagState &state, const AugLagParam &params,
                       const Eigen::VectorXd &x, Eigen::VectorXd &g) {
    auto hi = evaluate(x, g);
    auto newlambda_ = lambda_ + state.rho * hi;
    state.penalty += fabs(hi);
    state.feasible = state.feasible && fabs(hi) <= tol_;
    state.ICM = std::max(state.ICM, fabs(hi));
    lambda_ = clamp(params.lambda_min, params.lambda_max, newlambda_);
    return fabs(hi);
  }

  double forward(AugLagState &state, const AugLagParam &params,
                 const Eigen::VectorXd &x, Eigen::VectorXd &g) {
    auto res_tmp = evaluate(x, g);
    auto h = res_tmp + lambda_ / state.rho;
    auto L = 0.5 * state.rho * h * h;
    g = state.rho * h * g;
    return L;
  }
};

// f(x) <= 0
class LagInEqConstraint : public LagConstraintScalarBase {
private:
  double tol_ = 1e-6;
  double lambda_ = 0;

public:
  double update_lambda(AugLagState &state, const AugLagParam &params,
                       const Eigen::VectorXd &x, Eigen::VectorXd &g) {
    auto fci = evaluate(x, g);
    auto newlambda_ = lambda_ + state.rho * fci;
    state.penalty += fci > 0 ? fci : 0;
    state.feasible = state.feasible && fci <= tol_;
    state.ICM = std::max(state.ICM, fabs(std::max(fci, -lambda_ / state.rho)));
    lambda_ = clamp(0, params.lambda_max, newlambda_);
    return fci > 0 ? fci : 0;
  }

  double forward(AugLagState &state, const AugLagParam &params,
                 const Eigen::VectorXd &x, Eigen::VectorXd &g) {
    auto res_tmp = evaluate(x, g);
    auto h = res_tmp + lambda_ / state.rho;
    double L = 0;
    if (h > 0) {
      L = 0.5 * state.rho * h * h;
      g = state.rho * h * g;
    } else {
      g.fill(0);
    }
    return L;
  }
};

enum class AugLagReturnType {
  CONVERGENCE = 0,
  STOP,
  CANCELED,
  // ERROR
  ERR_UNKNOWN = -1024,
  ERR_INVALIDPARAMS,
  ERR_INVALIDLOCALSOLVER,
  ERR_MAXIMUMITERATION,
  ERR_MAXIMUMTIMELIMIT,
  ERR_LOCALERROR,
};

class AugLagSolver : public BaseSolver {
private:
  std::vector<std::shared_ptr<LagConstraintBase>> constraints_;
  std::shared_ptr<BaseSolver> local_solver_;
  evaluate_func_type objective_;
  void *data_ = nullptr;
  int local_ret_ = 0;
  AugLagState state_;
  Eigen::VectorXd g_tmp_; // cache data

  print_auglag_progress_func_type print_progress_ =
      default_auglag_print_progress;

  static double default_auglag_evaluate_function(void *self,
                                                 const Eigen::VectorXd &x,
                                                 Eigen::VectorXd &g) {
    auto ptr = static_cast<AugLagSolver *>(self);
    auto L = ptr->objective_(ptr->data_, x, g);
    for (auto &cons : ptr->constraints_) {
      L += cons->forward(ptr->state_, ptr->param, x, ptr->g_tmp_);
      g += ptr->g_tmp_;
    }
    return L;
  }

public:
  AugLagParam param;
  AugLagSolver() = default;

  void add_constraint(const std::shared_ptr<LagConstraintBase> &constraint) {
    constraints_.push_back(constraint);
  }

  void set_local_solver(const std::shared_ptr<BaseSolver> &solver) {
    local_solver_ = solver;
  }

  void set_objective(const evaluate_func_type &objective,
                     void *instance) override {
    objective_ = objective;
    data_ = instance;
  }

  void set_progress(const print_auglag_progress_func_type &func) {
    print_progress_ = func;
  }

  int get_local_ret() const { return local_ret_; }

  int minimize(Eigen::VectorXd &x, double &minf) override {
    if (!local_solver_) {
      return (int)AugLagReturnType::ERR_LOCALERROR;
    }
    if (!objective_) {
      return (int)AugLagReturnType::ERR_INVALIDPARAMS;
    }
    auto start = std::chrono::steady_clock::now();
    /* Set the minimization parameters */

    /* magic parameters from Birgin & Martinez */
    state_.ICM = HUGE_VAL;
    state_.penalty = 0;
    state_.feasible = true;
    state_.rho = 1.0;

    AugLagState minf_state;
    minf_state.penalty = HUGE_VAL;
    minf_state.feasible = false;

    minf = HUGE_VAL;
    g_tmp_.resize(x.size());

    /* starting rho suggested by B & M */
    if (!constraints_.empty()) {
      double con2 = 0;
      // ++ *(d.stop->nevals_p);
      auto fcur = objective_(data_, x, g_tmp_);
      // if (nlopt_stop_forced(stop)) {
      //   ret = NLOPT_FORCED_STOP; goto done; }
      state_.penalty = 0;
      state_.feasible = true;
      for (auto &cons : constraints_) {
        auto hi = cons->update_lambda(state_, param, x, g_tmp_);
        con2 += hi * hi;
      }

      minf = fcur;
      minf_state.penalty = state_.penalty;
      minf_state.feasible = state_.feasible;
      state_.rho = (con2 > 0) ? clamp(1e-6, 10, 2.0 * fabs(minf) / con2) : 10;
    } else {
      state_.rho = 1; /* whatever, doesn't matter */
    }

    AugLagReturnType ret = AugLagReturnType::CONVERGENCE;
    if (print_progress_) {
      int monitor_ret = print_progress_(state_, minf, x, g_tmp_);
      if (monitor_ret < 0) {
        ret = AugLagReturnType::CANCELED;
        return (int)ret;
      }
    }

    double fcur = HUGE_VAL;

    Eigen::VectorXd x_prev = x;

    local_solver_->set_objective(default_auglag_evaluate_function, this);

    state_.iteration = 0;
    while (true) {
      double prev_ICM = state_.ICM;
      // optimize;
      /* Start minimization */
      local_ret_ = local_solver_->minimize(x, fcur);
      // break;
      // stop func
      if (local_ret_ < -128) {
        ret = AugLagReturnType::ERR_LOCALERROR;
        break;
      }

      state_.ICM = 0;
      state_.penalty = 0;
      state_.feasible = true;

      for (auto &cons : constraints_) {
        cons->update_lambda(state_, param, x, g_tmp_);
      }

      if (state_.ICM > param.tau * prev_ICM) {
        state_.rho *= param.gam;
      }

      auto fcur = objective_(data_, x, g_tmp_);

      state_.iteration++;

      if (print_progress_) {
        int monitor_ret = print_progress_(state_, fcur, x, g_tmp_);
        if (monitor_ret < 0) {
          ret = AugLagReturnType::CANCELED;
          return (int)ret;
        }
      }

      if ((state_.feasible &&
           (!minf_state.feasible || state_.penalty < minf_state.penalty ||
            fcur < minf)) ||
          (!minf_state.feasible && state_.penalty < minf_state.penalty)) {
        ret = AugLagReturnType::CONVERGENCE;
        if (state_.feasible) {
          const bool fIsOk =
              (fabs(minf - fcur) / std::max(1.0, fabs(fcur)) < param.delta);
          if (fIsOk) {
            const auto gnorm_inf = g_tmp_.cwiseAbs().maxCoeff();
            const auto xnorm_inf = x.cwiseAbs().maxCoeff();
            if (gnorm_inf / std::max(1.0, xnorm_inf) < param.g_epsilon) {
              ret = AugLagReturnType::CONVERGENCE;
            } else {
              ret = AugLagReturnType::STOP;
            }
            break;
          }
        }
        minf = fcur;
        minf_state.penalty = state_.penalty;
        minf_state.feasible = state_.feasible;
        x_prev = x;
        if ((int)ret != 0)
          break;
      }
      if (state_.ICM == 0) {
        ret = AugLagReturnType::STOP;
        break;
      }

      if (param.max_iterations > 0 && state_.iteration > param.max_iterations) {
        ret = AugLagReturnType::ERR_MAXIMUMITERATION;
        break;
      }

      if (param.max_time_seconds > 0) {
        auto end = std::chrono::steady_clock::now();
        auto elapsed =
            std::chrono::duration_cast<std::chrono::seconds>(end - start);
        state_.time_seconds = elapsed.count();
        if (state_.time_seconds > param.max_time_seconds) {
          ret = AugLagReturnType::ERR_MAXIMUMTIMELIMIT;
          break;
        }
      }
    }
    return (int)ret;
  }
};
} // namespace GOptSolver